package main

import (
	"fmt"
	"math"
	"time"

	"github.com/montanaflynn/stats"
	"github.com/sajari/regression"
)

type LoadForecaster struct {
	histCoeffs    []float64   // 历史模型系数
	realTimeModel *regression.Regression
	lastRealValue float64
}

func NewLoadForecaster() *LoadForecaster {
	return &LoadForecaster{
		realTimeModel: new(regression.Regression),
	}
}

// 训练历史数据模型 (使用线性回归简化实现)
func (lf *LoadForecaster) TrainHistorical(historicalData []float64) {
	var xVals []float64
	for i := range historicalData {
		xVals = append(xVals, float64(i))
	}

	r := new(regression.Regression)
	r.SetObserved("负载")
	r.SetVar(0, "时间")

	for i, y := range historicalData {
		r.Train(regression.DataPoint(y, []float64{float64(i)}))
	}

	r.Run()
	lf.histCoeffs = []float64{r.Coeff(0), r.Coeff(1)}
}

// 更新实时数据模型
func (lf *LoadForecaster) UpdateRealTime(realTimeData []float64) {
	lf.realTimeModel = new(regression.Regression)
	lf.realTimeModel.SetObserved("实时负载")
	lf.realTimeModel.SetVar(0, "前值")

	for i := 1; i < len(realTimeData); i++ {
		lf.realTimeModel.Train(regression.DataPoint(realTimeData[i], []float64{realTimeData[i-1]}))
	}

	lf.realTimeModel.Run()
	lf.lastRealValue = realTimeData[len(realTimeData)-1]
}

// 组合预测
func (lf *LoadForecaster) Predict(steps int) []float64 {
	histPred := make([]float64, steps)
	lastTime := float64(len(lf.histCoeffs)) // 假设最后时间点

	// 历史趋势预测
	for i := range histPred {
		histPred[i] = lf.histCoeffs[0] + lf.histCoeffs[1]*(lastTime+float64(i))
	}

	// 实时数据预测
	rtPred := make([]float64, steps)
	current := lf.lastRealValue
	for i := range rtPred {
		pred, _ := lf.realTimeModel.Predict([]float64{current})
		rtPred[i] = pred
		current = pred
	}

	// 组合预测 (70%历史 + 30%实时)
	combined := make([]float64, steps)
	for i := range combined {
		combined[i] = 0.7*histPred[i] + 0.3*rtPred[i]
	}

	return combined
}

func main() {
	// 示例数据
	historical := make([]float64, 30*24)
	for i := range historical {
		historical[i] = 500 + 50*math.Sin(float64(i)/24*2*math.Pi)
	}

	realtime := []float64{480, 490, 510, 520, 515, 530, 525, 540}

	// 创建预测器
	forecaster := NewLoadForecaster()
	forecaster.TrainHistorical(historical)
	forecaster.UpdateRealTime(realtime)

	// 预测未来24个点
	prediction := forecaster.Predict(24)
	fmt.Println("预测结果:", prediction)
}